import os
import torch
import torchvision.transforms as transforms
from torchvision.models import resnet50
from PIL import Image
import faiss
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from tqdm import tqdm


class LocalImageSearch:
    def __init__(self, model_path=None, use_gpu=False):
        """初始化图像检索系统"""
        # 设备配置
        self.device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")

        # 加载模型
        self.model = self._load_model(model_path)

        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        # 初始化FAISS索引
        self.feature_dim = 2048
        self.index = faiss.IndexFlatL2(self.feature_dim)
        self.image_paths = []

    def _load_model(self, model_path):
        """加载ResNet-50模型，用于特征提取"""
        # 加载预训练模型
        model = resnet50(pretrained=True)
        # 移除最后一层分类层，保留特征提取部分
        model = torch.nn.Sequential(*list(model.children())[:-1])

        # 如果有微调模型，加载微调权重
        if model_path and os.path.exists(model_path):
            print(f"加载微调模型: {model_path}")
            pretrained_dict = torch.load(model_path, map_location=self.device)
            model_dict = model.state_dict()
            # 过滤不需要的参数
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            model_dict.update(pretrained_dict)
            model.load_state_dict(model_dict)

        model.to(self.device)
        model.eval()  # 设置为评估模式
        return model

    def extract_feature(self, image_path):
        """提取单张图片的特征向量"""
        try:
            # 打开图片并转换为RGB
            img = Image.open(image_path).convert("RGB")
            # 预处理
            img_tensor = self.transform(img).unsqueeze(0).to(self.device)

            # 提取特征
            with torch.no_grad():  # 禁用梯度计算，加速并节省内存
                feature = self.model(img_tensor)
                feature = feature.squeeze().cpu().numpy()  # 转换为numpy数组
                return feature.reshape(1, -1)  # 调整形状以适应FAISS
        except Exception as e:
            print(f"提取特征失败 {image_path}: {str(e)}")
            return None

    def build_index_from_directory(self, image_dir):
        """从目录构建图像索引"""
        print(f"从目录构建索引: {image_dir}")
        self.image_paths = []
        self.index.reset()  # 重置索引

        # 获取所有图片文件
        image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
        image_files = []

        for root, _, files in os.walk(image_dir):
            for file in files:
                if file.lower().endswith(image_extensions):
                    image_files.append(os.path.join(root, file))

        # 批量提取特征并添加到索引
        for img_path in tqdm(image_files, desc="构建索引"):
            feature = self.extract_feature(img_path)
            if feature is not None:
                self.index.add(feature)
                self.image_paths.append(img_path)

        print(f"索引构建完成，共处理 {len(self.image_paths)} 张图片")
        return len(self.image_paths)

    def search_similar(self, query_image_path, top_k=5):
        """搜索相似图片"""
        # 提取查询图片特征
        query_feature = self.extract_feature(query_image_path)
        if query_feature is None:
            return None

        # 搜索相似图片
        distances, indices = self.index.search(query_feature, top_k)

        # 整理结果
        results = []
        for i in range(len(indices[0])):
            idx = indices[0][i]
            if idx < 0 or idx >= len(self.image_paths):
                continue

            img_path = self.image_paths[idx]
            distance = distances[0][i]
            # 归一化相似度 (0-1之间，值越大越相似)
            similarity = 1.0 / (1.0 + distance / 1000)

            results.append({
                "image_path": img_path,
                "similarity": round(similarity, 3),
                "distance": round(distance, 2)
            })

        return results

    def save_index(self, index_dir):
        """保存索引以便后续使用"""
        os.makedirs(index_dir, exist_ok=True)
        # 保存FAISS索引
        faiss.write_index(self.index, os.path.join(index_dir, "image_index.index"))
        # 保存图片路径列表
        with open(os.path.join(index_dir, "image_paths.txt"), "w", encoding="utf-8") as f:
            for path in self.image_paths:
                f.write(f"{path}\n")
        print(f"索引已保存至 {index_dir}")

    def load_index(self, index_dir):
        """加载已保存的索引"""
        try:
            # 加载FAISS索引
            self.index = faiss.read_index(os.path.join(index_dir, "image_index.index"))
            # 加载图片路径列表
            self.image_paths = []
            with open(os.path.join(index_dir, "image_paths.txt"), "r", encoding="utf-8") as f:
                for line in f:
                    path = line.strip()
                    if path and os.path.exists(path):
                        self.image_paths.append(path)
            print(f"索引已加载，共 {len(self.image_paths)} 张图片")
            return True
        except Exception as e:
            print(f"加载索引失败: {str(e)}")
            return False


def visualize_results(query_path, results):
    """可视化查询结果"""
    if not results:
        print("没有找到相似图片")
        return

    # 创建画布
    plt.figure(figsize=(15, 5))

    # 显示查询图片
    plt.subplot(1, len(results) + 1, 1)
    try:
        img = mpimg.imread(query_path)
        plt.imshow(img)
        plt.title("查询图片")
        plt.axis('off')
    except Exception as e:
        print(f"无法显示查询图片: {str(e)}")

    # 显示结果图片
    for i, result in enumerate(results, 2):
        plt.subplot(1, len(results) + 1, i)
        try:
            img = mpimg.imread(result["image_path"])
            plt.imshow(img)
            plt.title(f"相似度: {result['similarity']}")
            plt.axis('off')
        except Exception as e:
            print(f"无法显示图片 {result['image_path']}: {str(e)}")

    plt.tight_layout()
    plt.show()


def main():
    # 配置参数
    IMAGE_DIR = "../shop/product_images"  # 商品图片目录
    INDEX_DIR = "../suoyin/image_index"  # 索引保存目录
    QUERY_IMAGE = "../shop/query_images/2362.jpg"  # 查询图片路径
    TOP_K = 5  # 返回最相似的前5张图片
    USE_GPU = False  # 是否使用GPU

    # 初始化检索系统
    search_engine = LocalImageSearch(use_gpu=USE_GPU)

    # 检查是否有已保存的索引
    if os.path.exists(INDEX_DIR) and len(os.listdir(INDEX_DIR)) > 0:
        search_engine.load_index(INDEX_DIR)
    else:
        # 构建新索引
        if not os.path.exists(IMAGE_DIR):
            print(f"商品图片目录不存在: {IMAGE_DIR}")
            print("请创建该目录并放入商品图片后重试")
            return

        search_engine.build_index_from_directory(IMAGE_DIR)
        # 保存索引以便下次使用
        search_engine.save_index(INDEX_DIR)

    # 检查查询图片是否存在
    if not os.path.exists(QUERY_IMAGE):
        print(f"查询图片不存在: {QUERY_IMAGE}")
        print("请将查询图片命名为query.jpg并放在当前目录")
        return

    # 执行搜索
    print(f"正在搜索与 {QUERY_IMAGE} 相似的图片...")
    results = search_engine.search_similar(QUERY_IMAGE, top_k=TOP_K)

    # 显示结果
    if results:
        print("\n搜索结果:")
        print(results)
        print("\n搜索结果:")
        for i, result in enumerate(results, 1):
            print(f"{i}. {result['image_path']} - 相似度: {result['similarity']}")

        # 可视化结果
        visualize_results(QUERY_IMAGE, results)
    else:
        print("搜索失败，未找到相似图片")


if __name__ == "__main__":
    main()
